from typing import Dict, Callable, Any
from .schemas import ToolDescription, ToolResult


class ToolManager:
    def __init__(self):
        self._tools: Dict[str, Callable] = {}
        self._descriptions: Dict[str, ToolDescription] = {}

    def register_tool(self, func: Callable, description: str = None):
        tool_name = func.__name__
        self._tools[tool_name] = func

        # 自动提取参数信息
        params = {}
        for name, param in func.__annotations__.items():
            if name != 'return':
                params[name] = {"type": str(param.__name__)}

        self._descriptions[tool_name] = ToolDescription(
            name=tool_name,
            description=description or func.__doc__ or "",
            parameters=params,
            return_type=func.__annotations__.get('return', 'Any').__name__
        )
        return func

    def list_tools(self) -> Dict[str, ToolDescription]:
        return self._descriptions

    async def invoke_tool(self, tool_name: str, parameters: Dict[str, Any]) -> ToolResult:
        if tool_name not in self._tools:
            return ToolResult(success=False, error=f"Tool {tool_name} not found")

        try:
            result = await self._tools[tool_name](**parameters)
            return ToolResult(success=True, output=result)
        except Exception as e:
            return ToolResult(success=False, error=str(e))


# 工具装饰器
def tool(description: str = None):
    def decorator(func):
        tool_manager.register_tool(func, description)
        return func

    return decorator


# 全局工具管理器实例
tool_manager = ToolManager()